import pandas as pd
import numpy as np
from pyexpat import features
from matplotlib import pyplot as plt
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
import seaborn as sns
from sklearn.preprocessing import MinMaxScaler, LabelEncoder


#定义函数
# def data_preprocessing(path):
#     data = pd.read_csv(path)
#     #删除Over18的列  全为Y  StandardHours 全为80
#     data = data.drop(columns=['DistanceFromHome','EmployeeNumber','Over18','StandardHours','Gender'])
#     data = pd.get_dummies(data, columns=['BusinessTravel','Department','EducationField','JobRole', 'MaritalStatus', 'OverTime'])
#     #增加一个新列   总年份除以总公司数
#     data.groupby('Education')['Attrition'].value_counts()
#     # data['year_company'] =data['NumCompaniesWorked']/data['TotalWorkingYears']
#     print(data.head(10).to_string())
#     # print(data.info())
#     return data
def data_preprocessing(path):
    data = pd.read_csv(path)
    # print(data.head(5).to_string())
    data = data.drop(columns=['Over18','Gender','Department','StandardHours','PerformanceRating'])
    # print(data.head(10).to_string())
    categorical_columns = data.select_dtypes(include=['object']).columns.tolist()
    # print(categorical_columns)
    for i in categorical_columns:
        le = LabelEncoder()
        data[i] = le.fit_transform(data[i])
    # print(data.head(10).to_string())
    # print(data.info())
    x = data.iloc[:,1:]
    y = data['Attrition']
    # print(x)
    # print(y)
    x_train,x_test,y_train,y_test = train_test_split(x,y,test_size=0.2,random_state=25,stratify=y)
    model = RandomForestClassifier()
    model.fit(x_train,y_train)


    #随机森林
    # feature_importance = model.feature_importances_
    # sorted_indices = feature_importance.argsort()[1::]
    # feature_names = x.columns
    # print(type(sorted_indices))
    # for index in sorted_indices:
        # print(feature_importance[index])


    #可视化展示特征的重要性
    # 绘制特征重要性条形图
    # plt.figure(figsize=(10, 8))
    # sns.barplot(x=feature_importance[sorted_indices], y=feature_names[sorted_indices])
    # plt.title('Feature Importance')
    # plt.xlabel('Importance')
    # plt.ylabel('Features')
    # plt.show()
    return data


if __name__ == '__main__':
    data_preprocessing(r"../../data/raw/train.csv")